#!/usr/bin/env python3
"""
Minimal trajectory visualization helper.

Currently generates KDE contour figures for:
1. Ego start/end distributions in the (x, z) plane.
2. Opponent start/end distributions in the (x, z) plane.
3. Opponent goal distribution (x, y) and ego endpoint distribution (x, y).
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Dict, Tuple

import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401


DEFAULT_DATASET_DIR = Path(__file__).resolve().parent / "vaedataset"


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate contour visualizations from a VAE dataset (.npz)."
    )
    parser.add_argument(
        "-i",
        "--input",
        default=str(DEFAULT_DATASET_DIR),
        help=(
            "Path to the .npz dataset or a directory containing datasets "
            f"(default: {DEFAULT_DATASET_DIR})."
        ),
    )
    parser.add_argument(
        "-o",
        "--output-dir",
        default=str(DEFAULT_DATASET_DIR),
        help=(
            "Directory to store generated figures "
            f"(default: {DEFAULT_DATASET_DIR})."
        ),
    )
    parser.add_argument(
        "--grid-size",
        type=int,
        default=200,
        help="Grid resolution for KDE contours (default: 200).",
    )
    return parser.parse_args()


def _load_dataset(dataset_path: Path) -> Dict[str, np.ndarray]:
    data = np.load(dataset_path)
    required = ["ego_trajectories", "opponent_trajectories", "opponent_goals"]
    missing = [k for k in required if k not in data.files]
    if missing:
        raise KeyError(f"Missing keys in dataset: {missing}")
    return data


def _reshape_trajectories(arr: np.ndarray, window: int) -> np.ndarray:
    data = arr.reshape(-1, window, 6)
    return data[:, :, :3]  # positions only


def _prepare_output_dir(base_dir: Path, dataset_path: Path) -> Path:
    base_dir.mkdir(parents=True, exist_ok=True)
    target_dir = base_dir / dataset_path.stem
    target_dir.mkdir(parents=True, exist_ok=True)
    return target_dir


def _resolve_dataset_path(raw_path: Path) -> Path:
    candidate = raw_path.expanduser().resolve()
    if candidate.is_file():
        return candidate
    if candidate.is_dir():
        npz_files = sorted(candidate.glob("*.npz"))
        if not npz_files:
            raise FileNotFoundError(f"No .npz files found under {candidate}")
        latest = max(npz_files, key=lambda p: p.stat().st_mtime)
        print(f"Directory provided, using latest dataset: {latest.name}")
        return latest
    if candidate.suffix != ".npz":
        raise FileNotFoundError(f"Dataset not found: {candidate}")
    raise FileNotFoundError(f"Dataset not found: {candidate}")


def _silverman_bandwidth(data: np.ndarray) -> float:
    std = np.std(data, ddof=1)
    if std == 0:
        return 1.0
    q75, q25 = np.percentile(data, [75, 25])
    iqr = q75 - q25
    sigma = min(std, iqr / 1.34) if iqr > 0 else std
    return 0.9 * sigma * (data.size ** (-0.2))


def _kde_density(
    x: np.ndarray,
    y: np.ndarray,
    grid_x: np.ndarray,
    grid_y: np.ndarray,
) -> np.ndarray:
    bw_x = max(_silverman_bandwidth(x), 1e-3)
    bw_y = max(_silverman_bandwidth(y), 1e-3)
    dx = (grid_x[..., None] - x[None, None, :]) / bw_x
    dy = (grid_y[..., None] - y[None, None, :]) / bw_y
    norm = 1.0 / (2 * np.pi * bw_x * bw_y)
    densities = np.exp(-0.5 * (dx**2 + dy**2)) * norm
    return densities.mean(axis=2)


def _plot_kde_contour(
    x: np.ndarray,
    y: np.ndarray,
    title: str,
    xlabel: str,
    ylabel: str,
    cmap: str,
    output_path: Path,
    grid_size: int,
) -> None:
    if x.size == 0 or y.size == 0:
        return
    xmin, xmax = float(x.min()), float(x.max())
    ymin, ymax = float(y.min()), float(y.max())
    x_margin = 0.1 * max(1e-3, xmax - xmin)
    y_margin = 0.1 * max(1e-3, ymax - ymin)
    xs = np.linspace(xmin - x_margin, xmax + x_margin, grid_size)
    ys = np.linspace(ymin - y_margin, ymax + y_margin, grid_size)
    grid_x, grid_y = np.meshgrid(xs, ys)
    density = _kde_density(x, y, grid_x, grid_y)

    fig, ax = plt.subplots(figsize=(6.5, 5.5))
    cf = ax.contourf(grid_x, grid_y, density, levels=15, cmap=cmap)
    cs = ax.contour(grid_x, grid_y, density, levels=10, colors="white", linewidths=0.5)
    ax.clabel(cs, inline=True, fontsize=7, fmt="%.2f")
    fig.colorbar(cf, ax=ax, label="Density")
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    fig.tight_layout()
    fig.savefig(output_path, dpi=300)
    plt.close(fig)


def main() -> None:
    args = _parse_args()
    dataset_path = _resolve_dataset_path(Path(args.input))

    output_dir = _prepare_output_dir(Path(args.output_dir), dataset_path)
    data = _load_dataset(dataset_path)

    window = int(data["observation_window"]) if "observation_window" in data.files else (
        data["ego_trajectories"].shape[1] // 6
    )
    ego_positions = _reshape_trajectories(data["ego_trajectories"], window)
    opponent_positions = _reshape_trajectories(data["opponent_trajectories"], window)
    ego_start = ego_positions[:, 0, :]
    ego_end = ego_positions[:, -1, :]
    opp_start = opponent_positions[:, 0, :]
    opp_end = opponent_positions[:, -1, :]

    print(f"Loaded dataset: {dataset_path.name}")
    print(f"Samples: {ego_positions.shape[0]}, window: {window}")
    print(f"Saving figures to: {output_dir}")

    _plot_kde_contour(
        x=ego_start[:, 0],
        y=ego_start[:, 2],
        title="Ego Start KDE (x-z)",
        xlabel="Start x (m)",
        ylabel="Start z (m)",
        cmap="magma",
        output_path=output_dir / "ego_start_kde_xz.png",
        grid_size=args.grid_size,
    )
    _plot_kde_contour(
        x=ego_end[:, 0],
        y=ego_end[:, 2],
        title="Ego End KDE (x-z)",
        xlabel="End x (m)",
        ylabel="End z (m)",
        cmap="plasma",
        output_path=output_dir / "ego_end_kde_xz.png",
        grid_size=args.grid_size,
    )

    _plot_kde_contour(
        x=opp_start[:, 0],
        y=opp_start[:, 2],
        title="Opponent Start KDE (x-z)",
        xlabel="Start x (m)",
        ylabel="Start z (m)",
        cmap="cividis",
        output_path=output_dir / "opponent_start_kde_xz.png",
        grid_size=args.grid_size,
    )
    _plot_kde_contour(
        x=opp_end[:, 0],
        y=opp_end[:, 2],
        title="Opponent End KDE (x-z)",
        xlabel="End x (m)",
        ylabel="End z (m)",
        cmap="inferno",
        output_path=output_dir / "opponent_end_kde_xz.png",
        grid_size=args.grid_size,
    )

    goals = data["opponent_goals"]
    _plot_kde_contour(
        x=goals[:, 0],
        y=goals[:, 1],
        title="Opponent Goal KDE Contour (x-y)",
        xlabel="Goal x (m)",
        ylabel="Goal y (m)",
        cmap="viridis",
        output_path=output_dir / "opponent_goal_xy_contour.png",
        grid_size=args.grid_size,
    )

    _plot_kde_contour(
        x=ego_end[:, 0],
        y=ego_end[:, 1],
        title="Ego End KDE Contour (x-y)",
        xlabel="End x (m)",
        ylabel="End y (m)",
        cmap="Blues",
        output_path=output_dir / "ego_end_xy_contour.png",
        grid_size=args.grid_size,
    )

    print(
        "✅ Generated ego/opponent start-end KDEs (xz) and "
        "ego/opponent XY contour maps."
    )


if __name__ == "__main__":
    main()
